#ifndef TACORES_PRIM_H
#define TACORES_PRIM_H

using namespace std;

#include "Graph.h"
#include <vector>
#include <queue>
#include <limits>

// O((V+E)*log V)
class Prim
{
public:
	Prim(Graph& g)
		: g (g)
	{
	}

	void computeMST(int s)
	{
		initialize();
		selectedEdge[s] = 0;
		priority_queue<P, vector<P>, greater<P> > pq;

		int size = g.numVertices();
		for(int v = 0; v < size; ++v)
		{
			pq.push(P(selectedEdge[v],v));
		}
		contain.assign(size, true);

		while(pq.empty() == false)
		{
			P p = pq.top(); pq.pop();
			int u = p.second;
			if(selectedEdge[u] < p.first)
			{
				continue;
			}
			contain[u] = false;

			VertexList::const_iterator end = g.end(u);
			for(VertexList::const_iterator it = g.begin(u); it != end; ++it)
			{
				int v = it->first;
				int w = it->second;
				if(contain[v])
				{
					if(w < selectedEdge[v])
					{
						pred[v] = u;
						selectedEdge[v] = w;
						pq.push(P(w,v));
						contain[v] = true;
					}
				}
			}
		}
	}

	vector<int> getSelectedEdge() { return selectedEdge; }
	vector<int> getPredecessor() { return pred; }

private:
	void initialize()
	{
		int size = g.numVertices();
		selectedEdge.assign(size, numeric_limits<int>::max());
		pred.assign(size, -1);
	}

	Graph& g;
	vector<int> selectedEdge;
	vector<int> pred;
	vector<bool> contain;

	typedef pair<int,int> P;
};
typedef Prim PRIM;

#endif

